# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: taekwon

import datetime
import functools
import logging

import os

from absl import app, flags

from coin.base.datetime_util import to_datetime

import experimental.prophet as prophet
from experimental.prophet.craps.universe import universe_all
from coin.strategy.interval.research import features
from coin.strategy.interval.model import LinearModel


def matrix(ref_symbol_nicks, date):
  vars = []

  # Timers, sampling points
  ts = prophet.global_timestamp()
  timer_1s = prophet.timer('1s')

  # Y-features
  prev_y_vars = []
  prev_y_vars = features.features(ref_symbol_nicks, date, True)

  # Features
  vars += features.features(ref_symbol_nicks, date)

  # TODO(taekwon): fix following later.
  # Sampler
  # with prophet.namespace('sampler'):
  #   samplers = features.samplers(target_coin, 'feature/')
  #   vars += samplers

  # Trigger
  # vars += [features.trigger(target_coin, symbol_nick, 'feature/')]

  # Long-term model theo price
  # with prophet.namespace('lt_model'):
  #   theo_ask0_p, theo_bid0_p = lt_model.theo_price(symbol_info, date, **lt_model_param)
  #   vars += [theo_ask0_p, theo_bid0_p]

  # Fill info
  # with prophet.namespace('fillinfo'):
  #   vars += y_features.multiple_edge_fillinfo(
  #       symbol_info, feed, theo_ask0_p, theo_bid0_p, trigger, order_latency, order_lifetime,
  #       taker_panalty_bp / 10000.)

  # 1 minute

  prev_ts = prophet.shift(ts, 1)
  trigger = (prev_ts % (60 * 10**9) == 0)
  prev_vars = prophet.shift(vars, 1)

  with prophet.control_if(trigger & (prev_ts != ts)):
    aggregator = prophet.scalar_value_aggregator([prev_ts, *prev_vars, *prev_y_vars])

  return aggregator


# TODO(taekwon): might be needed.
def make_y_inplace(df, y_windows=None):
  y_windows = y_windows or ['1s', '2s', '5s', '10s', '15s', '20s', '30s', '1m', '2m', '5m']
  # neg_ret_mask = (df['Y/sell_ret_Y%s' % w] < 0.)
  # df['Y/sell_edge_Y%s' % w][neg_ret_mask] += 1.


def train(df, products, alpha_trans_lst, y_feature, is_ratio=0.5):
  x_list = list()
  y = df[f'Y/{products[0]}/{y_feature}'].fillna(0)

  model = LinearModel()
  for alpha_trans in alpha_trans_lst:
    model.add(0.0, alpha_trans)

  model.print()
  y_pred, fit_scores = model.ridge_fit(df, y, is_ratio)
  model.print()
  return model, y_pred, fit_scores


def get_dump_root_dir():
  return '/tmp/%s/interval_data_new' % os.path.expanduser('~').replace('/home/', '')


def main(argv):
  FLAGS = flags.FLAGS
  assert FLAGS.machine

  ref_symbols = FLAGS.ref_symbols.split(',')
  target_coin = FLAGS.target_coin
  if target_coin == 'btc':
    total_symbols = sorted(
        set([
            'bitflyer-fxbtcjpy',
            'bitmex-xbtusd',
            'huobi-btcusd-quarter',
            'okex-btcusd-quarter',
            'okex-btcusd-this_week',
            'okex-btcusd-next_week',
            'okex-btcusd-perpetual',
            'binance-btcusdt',
            'gdax-btcusd',
        ] + ref_symbols))

  elif target_coin == 'eth':
    total_symbols = sorted(
        set([
            'bitmex-ethusd',
            'huobi-ethusd-quarter',
            'okex-ethusd-quarter',
            'okex-ethusd-this_week',
            'okex-ethusd-next_week',
            'okex-ethusd-perpetual',
            'binance-ethusdt',
            'gdax-ethusd',
        ] + ref_symbols))
  elif target_coin == 'test':
    total_symbols = sorted(
        set([
            'okex-btcusd-quarter',
            'okex-btcusd-this_week',
            'okex-btcusd-perpetual',
        ]))
  else:
    raise ValueError(target_coin)

  logging.info('Start Dumping...')

  # end_date = datetime.datetime.now().date() - datetime.timedelta(days=2)
  # begin_date = end_date - datetime.timedelta(days=30)
  begin_date_str, end_date_str = FLAGS.trading_dates.split('-')
  begin_date = datetime.datetime.strptime(begin_date_str, "%Y%m%d")
  end_date = datetime.datetime.strptime(end_date_str, "%Y%m%d")
  logging.info('%s  -  %s', str(begin_date), str(end_date))

  friday_filter = lambda d: (d.weekday() != 4)
  df = prophet.run_from_fastfeed_multiprocess(
      functools.partial(matrix, total_symbols),
      begin_date,
      end_date,
      machine=FLAGS.machine,
      max_workers=None,
      use_run_cache=False,
      # debug_executor=FLAGS.debug_executor,
      inject_date_to_graph_func=True,
      date_filter=friday_filter,
      duration_before='2h',
      duration_after='2h')
  replace_map = dict()
  # remove .shift1
  for column in df.columns:
    replace_map[column] = column.rsplit('.', 1)[0]
    if replace_map[column][0] == 'Y':
      replace_map[column] = replace_map[column].rsplit('.', 1)[0]
  df.rename(columns=replace_map, inplace=True)

  print(df)
  print(df.describe())

  if flags.FLAGS.dump:
    ts_for_dumping = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    # ts_for_dumping = '0'
    dump_root_dir = get_dump_root_dir()
    base_folder = f'{dump_root_dir}'
    folder_name = f'{base_folder}/dump_{ts_for_dumping}'
    os.system(f'mkdir -p {folder_name}')
    dump_file = folder_name + '/dump.pkl.gz'
    df.to_pickle(dump_file, compression='gzip')
    print(f'{folder_name}')
    os.system(f'mkdir -p {dump_root_dir}/recent')
    os.system(f'ln -sfn {folder_name} {dump_root_dir}/recent')


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_boolean('dump', True, 'Dump')

  flags.DEFINE_string('target_coin', 'btc', 'kind of coin. btc, eth')

  flags.DEFINE_string('ref_symbols', '', 'upbit-btckrw,kraken-btcusd-perpetual')

  flags.DEFINE_boolean('debug_executor', False, 'Debug executor')

  flags.DEFINE_string('trading_dates',
                      '20190715-20190815',
                      'StartDate-EndDate, ex)20190715-20190815')

  flags.DEFINE_string('machine',
                      'feed-01.ap-northeast-1.aws',
                      'Machine such as feed-01.ap-northeast-1.aws, feed-01.cn-hongkong.aliyun')

  app.run(main)
